import torch
import torch.nn as nn
import torch.optim as optim
from app.models.moisture_model import MoistureModel
from app.utils.data_loader import get_dataloader
from app.utils.logger import setup_logger
from config.settings import CHECKPOINT_DIR, EPOCHS, LR, BATCH_SIZE

logger = setup_logger(__name__)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MoistureModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)

train_loader = get_dataloader("train", "moisture", BATCH_SIZE)
val_loader = get_dataloader("val", "moisture", BATCH_SIZE, shuffle=False)

best_acc = 0.0
for epoch in range(EPOCHS):
    model.train()
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # 验证
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    acc = correct / total
    logger.info(f"Epoch {epoch + 1}/{EPOCHS}, Val Acc: {acc:.4f}")

    if acc > best_acc:
        best_acc = acc
        torch.save(model.state_dict(), CHECKPOINT_DIR / "moisture_best.pth")